import numpy as np
import tensorflow as tf
import t3f


def tt_svd(ph_input_weight, tt_shape, tt_rank, layer_name):
	# svd should be executed on cpu
	with tf.device('/cpu:0'):
		with tf.name_scope('truancated_svd' + '_' + layer_name):
			tt_rank = [1] + tt_rank + [1]

			# change ph_input_weight to matrix form
			conv_matrix = tf.reshape(tf.transpose(ph_input_weight, [0,2,1,3]), [np.prod(tt_shape[0]), np.prod(tt_shape[1])])

			# input weight to tt, max rank is the matrix full rank
			tt_input = t3f.to_tt_matrix(conv_matrix, tt_shape, np.min((np.prod(tt_shape[0]), np.prod(tt_shape[1]))))

			# output weight got by tt-round
			tt_output = t3f.round(tt_input, max_tt_rank = tt_rank)

		return list(tt_output._tt_cores)
